import requests, json

from ai.common import process_messages, process_content
from config import AI_KEY_QIAN_FAN
from ai.consts import PLATFORM_QIAN_FAN

# 可用模型名称
free_qianfan_model = [
    "ernie-speed-8k",
    "ernie-speed-128k",
    "ernie-lite-8k",
    # 这个模型太弱了，根据提示词输出不了json
    # "ernie-tiny-8k"
]

def get_result_by_qianfan(data, ai_info):
    url = ai_info['url']
    
    payload = json.dumps({
        "model": ai_info['model'],
        "messages": process_messages(data)
    })
    headers = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {ai_info["api_key"]}'
    }
    
    response = requests.request("POST", url, headers=headers, data=payload)
    
    content = response.json()['choices'][0]['message']['content']
    
    return process_content(content, ai_info)

def get_qianfan_result(data, model_name):
    qianfan_ai_info = {
        "api_key": AI_KEY_QIAN_FAN,
        "model": model_name,
        "url": "https://qianfan.baidubce.com/v2/chat/completions",
        "platform": PLATFORM_QIAN_FAN
    }
    return get_result_by_qianfan(data, qianfan_ai_info)
